"""
@Description :   有向图建立
@Author      :   Li Junjie 
@Time        :   2025/03/19 20:42:32
"""
import pandas as pd
import networkx as nx
import numpy as np
import matplotlib.pyplot as plt
from sklearn.decomposition import PCA
from sklearn.preprocessing import StandardScaler
from geopy.distance import geodesic
import re
import os
import matplotlib.colors as mcolors
import json
import ast
from matplotlib.lines import Line2D
from collections import defaultdict

print("Step 1")
# Step 1: Read GIS Grid Data and other data tables from Excel
file_path = r"data\盐津河流域1km网格GIS数据20240925-backup.xlsx"
# 读取地理基础信息
gis_data = pd.read_excel(file_path, sheet_name='地理基础信息', header=[0, 1, 2])

# 检查是否包含海拔列，如果有则添加，否则跳过
possible_elevation_columns = ['海拔', '海拔（m）', 'Elevation', 'Elevation (m)']
elevation_column = None

for col in possible_elevation_columns:
    if col in gis_data.columns:
        elevation_column = col
        break

if elevation_column:
    gis_data = gis_data.rename(columns={elevation_column: '海拔'})
else:
    print("Warning: No suitable elevation column found in the data.")

# 重命名其他列
gis_data.columns = ['流域网格编码', '单个网格流域面积（km2）', '经度（°\xb0\xb0）', '纬度（°\xb0\xb0）', '坡度（均值/°/\xb0/\xb0）', '坡向（均值/°/\xb0/\xb0）', '海拔']

# 读取气象信息（例如降雨量）
meteorological_data = pd.read_excel(file_path, sheet_name='不同年月气象信息', header=[0])
meteorological_data.columns = ['流域网格编码', '年份', '月份', '降雨量']

# 读取人口和 GDP 数据
population_gdp_data = pd.read_excel(file_path, sheet_name='不同年份人口经济', header=[0])
population_gdp_data.columns = ['流域网格编码', '年份', '人口', 'GDP']

# 读取土地利用数据
land_use_data = pd.read_excel(file_path, sheet_name='不同年份土地利用类型', header=[0, 1, 2])
land_use_data.columns = ['流域网格编码', '年份', '农田（km2）', '森林（km2）', '灌木（km2）', '草原（km2）', '水域（km2）',
                         '冰雪（km2）', '裸地（km2）', '不透水面（km2）', '湿地（km2）']

# 合并人口、GDP、降雨量数据到 gis_data 中
gis_data = gis_data.merge(population_gdp_data, on='流域网格编码', how='left')
mean_rainfall_data = meteorological_data.groupby('流域网格编码')['降雨量'].mean().reset_index()
gis_data = gis_data.merge(mean_rainfall_data, on='流域网格编码', how='left')

# 将人口、GDP、降雨量数据转换为数值类型
gis_data['人口'] = pd.to_numeric(gis_data['人口'], errors='coerce')
gis_data['GDP'] = pd.to_numeric(gis_data['GDP'], errors='coerce')
gis_data['降雨量'] = pd.to_numeric(gis_data['降雨量'], errors='coerce')

elevation_available = '海拔' in gis_data.columns
if elevation_available:
    print("Step 4: Create Graph Based on Elevation Relationships")
    # Step 4: 基于海拔关系创建图形
    G = nx.DiGraph()
    # 添加所有节点
    for index, row in gis_data.iterrows():
        # 提取土地利用类型的面积
        land_use = {
            '农田': row.get('农田（km2）', 0),
            '森林': row.get('森林（km2）', 0),
            '灌木': row.get('灌木（km2）', 0),
            '草原': row.get('草原（km2）', 0),
            '水域': row.get('水域（km2）', 0),
            '冰雪': row.get('冰雪（km2）', 0),
            '裸地': row.get('裸地（km2）', 0),
            '不透水面': row.get('不透水面（km2）', 0),
            '湿地': row.get('湿地（km2）', 0)
        }

        # 将土地利用、人口、GDP、降雨量作为节点特征
        G.add_node(str(row['流域网格编码']),  # 使用字符串格式的节点ID
                   latitude=row['纬度（°\xb0\xb0）'],
                   longitude=row['经度（°\xb0\xb0）'],
                   elevation=row['海拔'],
                   slope=row['坡度（均值/°/\xb0/\xb0）'],
                   aspect=row['坡向（均值/°/\xb0/\xb0）'],
                   population=row['人口'],
                   gdp=row['GDP'],
                   rainfall=row['降雨量'],
                   land_use=land_use)

    # 从 grid.jsonl 文件中读取相邻网格数据，并建立节点连边
    grid_file_path = r"data/result.jsonl"
    with open(grid_file_path, 'r', encoding='utf-8') as f:
        for i, line in enumerate(f):
            try:
                # 使用 ast.literal_eval 解析
                grid_data = json.loads(line)
            except (ValueError, SyntaxError) as e:
                print(f"Error parsing line: {line.strip()}")
                continue
            # 获取当前节点和邻居节点列表
            key, value = next(iter(grid_data.items()))
            neighbors = value
            current_node = str(key)  # 确保节点ID是字符串格式
            # 为每个相邻网格建立边，增加节点存在性检查
            if current_node not in G.nodes:
                print(f"Warning: Current node {current_node} not found in graph.")
                continue

            for neighbor in neighbors:
                neighbor = str(neighbor)  # 确保相邻节点ID是字符串格式
                if neighbor is not None and neighbor in G.nodes:
                    current_elevation = G.nodes[current_node]['elevation']
                    neighbor_elevation = G.nodes[neighbor]['elevation']

                    if neighbor_elevation > current_elevation:
                        # 先检查存不存在neighbor指向current_node的边 如果不存在则建立
                        if not G.has_edge(neighbor, current_node):
                            G.add_edge(neighbor, current_node)
                    elif neighbor_elevation < current_elevation:
                        if not G.has_edge(current_node, neighbor):
                            G.add_edge(current_node, neighbor)
    print(f"Total nodes in graph G: {G.number_of_nodes()}")
    print(f"Total edges in graph G: {G.number_of_edges()}")

    # 检查图是否为强连通图
    if nx.is_strongly_connected(G):
        print("The graph is connected.")
    else:
        print("The graph is NOT connected.")

    # 可视化网络
    pos = nx.spring_layout(G, seed=42)
    plt.figure(figsize=(12, 12))
    nx.draw(G, pos, with_labels=True, node_size=500, node_color="skyblue", font_size=10, edge_color='gray')
    plt.title("Network Graph Visualization Based on Elevation")
    plt.show()

nx.write_gexf(G, "data/network_graph.gexf")